#include "StreamLine.h"
#include <functional>
#include <algorithm>

#include "OCesium/src/Ellipsoid.cpp"
#include "OCesium/src/Matrix4.cpp"
#include "OCesium/src/Spherical.cpp"
#include "OCesium/src/TilingScheme.cpp"

#define ZScale 1.0

inline emscripten::val gen_typeArray(const std::vector<float> pointCoords)
{
    auto view = emscripten::typed_memory_view(pointCoords.size(), pointCoords.data());
    return emscripten::val::global("Float32Array").new_(view);
}

typedef std::function<Cesium::Cartesian3(int x, int y, int z)> GetWindUVFunc;

class StreamGrid
{
public:
    StreamGrid(int width1, int height1, int depth1, int _maxTakenNum)
        : width(width1), height(height1), maxTakenNum(_maxTakenNum)
    {
        gridData.resize(depth1);
        for (auto &g : gridData)
        {
            g.resize(width * height);
        }
    }
    bool isTaken(int x, int y, int z, bool ocupy)
    {
        auto &grid = gridData[z];
        int i = y * width + x;
        if (grid[i] > maxTakenNum)
        {
            return true;
        }
        if (ocupy)
        {
            grid[i] += 1;
        }
        return false;
    }

    int maxTakenNum;
    int width, height;
    std::vector<Uint8Vector> gridData;
};

class StreamlineTracer
{
public:
    StreamlineTracer(GetWindUVFunc _getWind, int maxTakenNum,
                     int w, int h, int d = 1) : width(w),
                                                height(h),
                                                depth(d),
                                                getWind(_getWind),
                                                grid(w, h, d, maxTakenNum)
    {
        MaxTakenNum = maxTakenNum;
        boundScale = {(float)width, (float)height, (float)depth};
    }

    Cesium::Cartesian3 mix(Cesium::Cartesian3 a, Cesium::Cartesian3 b, float ratio)
    {
        return {
            a.x * ratio + (1 - ratio) * b.x,
            a.y * ratio + (1 - ratio) * b.y,
            a.z * ratio + (1 - ratio) * b.z};
    }

    Cesium::Cartesian3 normalizedVectorField(Cesium::Cartesian3 pos)
    {
        int lx = floor(pos.x);
        int ly = floor(pos.y);
        int lz = floor(pos.z);
        int ux = ceil(pos.x);
        int uy = ceil(pos.y);
        int uz = ceil(pos.z);

        if (lx < 0)
            lx = ux;
        if (ux >= width)
            ux = lx;
        if (ly < 0)
            ly = uy;
        if (uy > height)
            uy = ly;
        if (lz < 0)
            lz = uz;
        if (uz >= depth)
            uz = lz;

        if (lx < 0 || lx >= width ||
            ly < 0 || ly >= height ||
            ly < 0 || lz >= depth ||
            uz < 0 || ux >= width ||
            uy < 0 || uy >= height ||
            uz < 0 || uz >= depth)
        {
            return Cesium::Cartesian3::ZERO;
        }

        Cesium::Cartesian3 tl = getWind(lx, ly, lz);
        Cesium::Cartesian3 tr = getWind(lx + 1, ly, lz);
        Cesium::Cartesian3 bl = getWind(lx, ly + 1, lz);
        Cesium::Cartesian3 br = getWind(lx + 1, ly + 1, lz);

        if (tl == Cesium::Cartesian3::ZERO ||
            tr == Cesium::Cartesian3::ZERO ||
            bl == Cesium::Cartesian3::ZERO ||
            br == Cesium::Cartesian3::ZERO)
        {
            return Cesium::Cartesian3::ZERO;
        }

        Cesium::Cartesian3 res = mix(mix(tl, tr, pos.x - lx), mix(bl, br, pos.x - lx), 1 - pos.y + ly);

        return res.normalize(res);
    }

    Cesium::Cartesian3 rk4(Cesium::Cartesian3 point)
    {
        point *= boundScale;

        Cesium::Cartesian3 tl = normalizedVectorField(point);

#if 1
        Cesium::Cartesian3 tr = normalizedVectorField(point + Cesium::Cartesian3(1, 0, 0));
        Cesium::Cartesian3 bl = normalizedVectorField(point + Cesium::Cartesian3(0, 1, 0));
        Cesium::Cartesian3 br = normalizedVectorField(point + Cesium::Cartesian3(1, 1, 0));

        Cesium::Cartesian2 f;
        f.x = point.x - floor(point.x);
        f.y = point.y - floor(point.y);

        Cesium::Cartesian3 res = mix(mix(tl, tr, f.x), mix(bl, br, f.x), f.y);
        tl = res;
#endif
        tl.z *= ZScale;

        return tl / boundScale * timeStep;
    }

    bool isTaken(Cesium::Cartesian3 candidate, bool ocupy = true)
    {
        if (outOfBound(candidate))
        {
            return true;
        }
        candidate *= boundScale;
        return grid.isTaken(candidate.x, candidate.y, candidate.z, ocupy);
    }

    bool outOfBound(Cesium::Cartesian3 pos)
    {
        return pos.x < 0.0 ||
               pos.x >= 1.0 ||
               pos.y < 0.0 ||
               pos.y >= 1.0 ||
               pos.z < 0.0 ||
               pos.z >= 1.0;
    }

    Cesium::Cartesian3 growByVelocity(Cesium::Cartesian3 pos, Cesium::Cartesian3 velocity)
    {

        Cesium::Cartesian3 candidate = pos + velocity;
        if (candidate.z < 0.0)
        {
            candidate.z = 0.0;
        }

        if (outOfBound(candidate))
        {
            return Cesium::Cartesian3::ZERO;
        }

        if (isTaken(candidate))
        {
            return Cesium::Cartesian3::ZERO;
        }

        return candidate;
    }

    Cesium::Cartesian3 growBackward(Cesium::Cartesian3 pos)
    {
        Cesium::Cartesian3 velocity = rk4(pos);
        if (velocity == Cesium::Cartesian3::ZERO)
        {
            return Cesium::Cartesian3::ZERO;
        }
        return growByVelocity(pos, -velocity);
    }

    Cesium::Cartesian3 growForward(Cesium::Cartesian3 pos)
    {
        Cesium::Cartesian3 velocity = rk4(pos);
        if (velocity == Cesium::Cartesian3::ZERO)
        {
            return Cesium::Cartesian3::ZERO;
        }
        return growByVelocity(pos, velocity);
    }

    Cesium::Cartesian3 tracerLine(Cesium::Cartesian3 start)
    {
        std::vector<Cesium::Cartesian3> line;
        Cesium::Cartesian3 current = start;
        while (true)
        {
            if (outOfBound(current))
            {
                break;
            }

            Cesium::Cartesian3 point = growForward(current);
            if (point != Cesium::Cartesian3::ZERO)
            {
                current = point;
                line.push_back(point);
            }
            else
            {
                break;
            }
        }

        current = start;
        while (true)
        {
            if (outOfBound(current))
            {
                break;
            }

            Cesium::Cartesian3 point = growBackward(current);
            if (point != Cesium::Cartesian3::ZERO)
            {
                current = point;
                line.insert(line.begin(), point);
            }
            else
            {
                break;
            }
        }
        if (line.size() > 2)
        {
            auto firstPoint = line[0];
            if (nextLevelStart)
            {
                nextLevelStart = false;
                lines.insert(lines.begin(), std::move(line));
            }
            else
            {
                lines.push_back(std::move(line));
            }
            return firstPoint;
        }

        return Cesium::Cartesian3::ZERO;
    }

    Cesium::Cartesian3 getNextStart()
    {
        float nextTimeStep = MaxTakenNum;

        for (; lastLineIndex < lines.size(); lastLineIndex++, currentLineCheckedIndex = lines[lastLineIndex - 1].size() * 1.0)
        {
            auto &line = lines[lastLineIndex];

            for (auto i = currentLineCheckedIndex; i < line.size(); i++)
            {
                auto p = line[i];
                Cesium::Cartesian3 velocity = normalizedVectorField(p * boundScale);
                if (velocity == Cesium::Cartesian3::ZERO)
                {
                    continue;
                }
                velocity /= boundScale;

                Cesium::Cartesian3 nextStart;
                nextStart.x = p.x - velocity.x * nextTimeStep;
                nextStart.y = p.y + velocity.y * nextTimeStep;
                nextStart.z = p.z;

                if (!isTaken(nextStart, false) && nextStart != Cesium::Cartesian3::ZERO && !outOfBound(nextStart))
                {
                    return nextStart;
                }

                nextStart.x = p.x + velocity.x * nextTimeStep;
                nextStart.y = p.y - velocity.y * nextTimeStep;
                nextStart.z = p.z;

                if (!isTaken(nextStart, false) && nextStart != Cesium::Cartesian3::ZERO && !outOfBound(nextStart))
                {
                    return nextStart;
                }

                nextStart.x = p.x + velocity.x * nextTimeStep;
                nextStart.y = p.y + velocity.y * nextTimeStep;
                nextStart.z = p.z;

                if (!isTaken(nextStart, false) && nextStart != Cesium::Cartesian3::ZERO && !outOfBound(nextStart))
                {
                    return nextStart;
                }
            }
        }

        return Cesium::Cartesian3::ZERO;
    }

    void run(int levels, std::vector<Cesium::Cartesian2> interestPoints)
    {
        for (float z = 1.0; z < levels; z++)
        {
            nextLevelStart = true;
            for (const auto &[x, y] : interestPoints)
            {
                Cesium::Cartesian3 start = {x, y, z / this->depth};
                while (true)
                {
                    Cesium::Cartesian3 validPoint = tracerLine(start);
                    if (validPoint == Cesium::Cartesian3::ZERO)
                    {
                        currentLineCheckedIndex++;
                    }
                    start = getNextStart();
                    if (start == Cesium::Cartesian3::ZERO)
                    {
                        break;
                    }
                    break;
                }
            }
        }
    }

    void run(int levels)
    {
        float size = 4.0;
        for (float z = 1.0; z < levels; z++)
        {
            nextLevelStart = true;
            for (float r = 1.0; r < size; r++)
            {
                for (float c = 1.0; c < size; c++)
                {
                    Cesium::Cartesian3 start = {c / size, r / size, z / this->depth};

                    while (true)
                    {
                        Cesium::Cartesian3 validPoint = tracerLine(start);
                        if (validPoint == Cesium::Cartesian3::ZERO)
                        {
                            currentLineCheckedIndex++;
                        }
                        start = getNextStart();
                        if (start == Cesium::Cartesian3::ZERO)
                        {
                            break;
                        }
                    }
                }
            }
        }
    }

    bool nextLevelStart;
    Cesium::Cartesian3 boundScale;
    size_t lastLineIndex = 0;
    size_t currentLineCheckedIndex = 0;
    StreamGrid grid;
    std::vector<std::vector<Cesium::Cartesian3>> lines;
    float timeStep = 0.5;
    GetWindUVFunc getWind;
    int width, height, depth = 1;
    int MaxTakenNum;
};

void StreamLine::setImage(const emscripten::val &jsImage)
{
    grid3dData = std::make_unique<Grid3dData>(jsImage);
}

std::tuple<int, int> StreamLine::trimBySpeed(std::vector<Cesium::Cartesian3> &line, float speed, float FrontSpeedClipFactor, float EndSpeedClipFactor)
{
    int width = grid3dData->width;
    int height = grid3dData->height;
    int depth = grid3dData->depth;

    int32 beg = 0;
    int32 end = 0;

    for (const auto &point : line)
    {
        auto uvw = grid3dData->getUV(point.x * width, point.y * height, point.z * depth);

        if (Cesium::Cartesian3::magnitudeSquared(uvw) > speed * FrontSpeedClipFactor)
        {
            break;
        }
        beg++;
    }

    for (auto it = line.rbegin(); it != line.rend(); ++it)
    {
        auto point = *it;
        auto uvw = grid3dData->getUV(point.x * width, point.y * height, point.z * depth);
        if (Cesium::Cartesian3::magnitudeSquared(uvw) > speed * EndSpeedClipFactor)
        {
            break;
        }
        end++;
    }

    end = (int32)line.size() - end;

    return std::make_tuple(beg, end);
};

std::vector<Cesium::Cartesian2> StreamLine::getInterestPoints(Cesium::Rectangle insterestExtent)
{
    std::vector<Cesium::Cartesian2> points;

    auto &rectangle = grid3dData->rectangle;

    auto intersectRectangle = Cesium::Rectangle::intersection(insterestExtent, rectangle);

    int startX = (intersectRectangle.west - rectangle.west) / grid3dData->dx;
    int startY = (intersectRectangle.south - rectangle.south) / grid3dData->dy;

    for (int r = startY; r < grid3dData->height; r++)
    {
        float lat = rectangle.south + r * grid3dData->dy;
        if (lat > intersectRectangle.north)
            break;
        for (int c = startX; c < grid3dData->width; c++)
        {
            float lon = rectangle.west + c * grid3dData->dx;
            if (lon > intersectRectangle.east)
                break;

            // std::cout << "lon:" << lon << ",lat:" << lat << "\n";
            float x = (lon - rectangle.west) / rectangle.getWidth();
            float y = (lat - rectangle.south) / rectangle.getHeight();
            points.push_back({x, y});
        }
    }

    return points;
}

const emscripten::val StreamLine::run(const emscripten::val &jsOptions)
{
    bool verbose = false;
    size_t bufferSize = 0;
    bool flat = true;
    float FrontSpeedClipFactor = 1.0;
    float EndSpeedClipFactor = 1.0;
    float filterMin = 25;
    int mininumPointPerline = 18;
    int maxTakenNum = 2;
    int maxinumLevels = 1000;

    bufferSize = jsOptions["bufferSize"].as<int>();

    std::optional<Cesium::Rectangle> interestExtent;

    if (!jsOptions["interestExtent"].isNull() && !jsOptions["interestExtent"].isUndefined())
    {
        interestExtent = js_getRectangle(jsOptions["interestExtent"]);
    }

    if (!jsOptions["verbose"].isNull() && !jsOptions["verbose"].isUndefined())
    {
        verbose = jsOptions["verbose"].as<bool>();
    }
    if (!jsOptions["flat"].isNull() && !jsOptions["flat"].isUndefined())
    {
        flat = jsOptions["flat"].as<bool>();
    }
    if (!jsOptions["maxinumLevels"].isNull() && !jsOptions["maxinumLevels"].isUndefined())
    {
        maxinumLevels = jsOptions["maxinumLevels"].as<int>();
    }

    if (!jsOptions["FrontSpeedClipFactor"].isNull() && !jsOptions["FrontSpeedClipFactor"].isUndefined())
    {
        FrontSpeedClipFactor = jsOptions["FrontSpeedClipFactor"].as<float>();
    }

    if (!jsOptions["EndSpeedClipFactor"].isNull() && !jsOptions["EndSpeedClipFactor"].isUndefined())
    {
        EndSpeedClipFactor = jsOptions["EndSpeedClipFactor"].as<float>();
    }

    if (!jsOptions["filterMin"].isNull() && !jsOptions["filterMin"].isUndefined())
    {
        filterMin = jsOptions["filterMin"].as<float>();
    }
    if (!jsOptions["mininumPointPerline"].isNull() && !jsOptions["mininumPointPerline"].isUndefined())
    {
        mininumPointPerline = jsOptions["mininumPointPerline"].as<int>();
    }
    if (!jsOptions["maxTakenNum"].isNull() && !jsOptions["maxTakenNum"].isUndefined())
    {
        maxTakenNum = jsOptions["maxTakenNum"].as<int>();
    }

    auto getUV = [&](int x, int y, int z)
    {
        return grid3dData->getUV(x, y, z);
    };

    float west = Cesium::CesiumMath::toDegrees(grid3dData->rectangle.west);
    float south = Cesium::CesiumMath::toDegrees(grid3dData->rectangle.south);
    float east = Cesium::CesiumMath::toDegrees(grid3dData->rectangle.east);
    float north = Cesium::CesiumMath::toDegrees(grid3dData->rectangle.north);

    std::vector<std::vector<Cesium::Cartesian3>> allLines;
    StreamlineTracer tracer(getUV, maxTakenNum, grid3dData->width, grid3dData->height, grid3dData->depth);

    size_t width = bufferSize;
    size_t height = bufferSize;
    size_t depth = 512;
    size_t comp = 3;

    if (flat)
    {
        if (interestExtent.has_value())
        {
            auto insterestPoints = getInterestPoints(interestExtent.value());
            tracer.run(std::min(grid3dData->depth, maxinumLevels), insterestPoints);
        }
        else
        {
            tracer.run(std::min(grid3dData->depth, maxinumLevels));
        }
        for (auto &line : tracer.lines)
        {
            std::vector<std::vector<Cesium::Cartesian3>> splitLines;
            splitLines.push_back(std::vector<Cesium::Cartesian3>());

            size_t count = 0;
            size_t simplify = mininumPointPerline;
            auto [beg, end] = trimBySpeed(line, filterMin, FrontSpeedClipFactor, EndSpeedClipFactor);
            if (end - beg > mininumPointPerline)
            {
                for (; beg != end; beg++)
                {
                    const auto &point = line[beg];

                    float longitude = west + point.x * (east - west);

                    float latitude;
                    if (grid3dData->flipY)
                    {
                        latitude = north - point.y * (north - south);
                    }
                    else
                    {
                        latitude = south + point.y * (north - south);
                    }

                    Cesium::Cartesian3 p(longitude, latitude, point.z);

                    if (splitLines[splitLines.size() - 1].size() == depth - 1)
                    {
                        splitLines[splitLines.size() - 1].push_back(p);

                        splitLines.push_back(std::vector<Cesium::Cartesian3>());
                        const auto &preLine = splitLines[splitLines.size() - 2];
                        splitLines[splitLines.size() - 1].push_back(preLine[preLine.size() - 2]);
                    }

                    splitLines[splitLines.size() - 1].push_back(p);
                }
            }

            for (auto &line : splitLines)
            {
                if (line.size() > 0)
                {
                    allLines.insert(allLines.begin(), std::move(line));
                    allLines.push_back(std::move(line));
                }
            }
        }
    }

    std::vector<float>
        arrayBufferView;
    arrayBufferView.resize(width * height * depth * comp);

    auto maxNumberLine = width * height;

    size_t lineId = 0;
    size_t maxPointNumOfLine = 0;

    for (const auto &line : allLines)
    {
        auto r = floor(lineId / width);
        auto c = lineId % width;
        if (line.size() > maxPointNumOfLine)
        {
            maxPointNumOfLine = line.size();
        }
        for (auto z = 0; z < line.size() && z < depth; z++)
        {
            auto &p = line[z];
            auto index = (z * width * height + r * width + c) * comp;
            auto lon = Cesium::CesiumMath::toRadians(p.x);
            auto lat = Cesium::CesiumMath::toRadians(p.y);
            arrayBufferView[index] = (lon - grid3dData->rectangle.west) / grid3dData->rectangle.getWidth();
            arrayBufferView[index + 1] = (lat - grid3dData->rectangle.south) / grid3dData->rectangle.getHeight();
            arrayBufferView[index + 2] = p.z;
        }

        lineId++;
        if (lineId == maxNumberLine)
        {
            break;
        }
    }

    if (verbose)
    {
        std::cout << grid3dData->toString()
                  << "\ttrack num:" << tracer.lines.size()
                  << "\tfilterMin:" << filterMin
                  << "\tfilter num:" << allLines.size()
                  << "\tmaxPointNumOfLine:" << maxPointNumOfLine;

        if (interestExtent.has_value())
        {
            std::cout << "\trectangle:" << interestExtent.value().toString();
        }
        std::cout << "\n";
    }

    return gen_typeArray(arrayBufferView);
}
